// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/common/memory/stream_arena.h>
#include <pollux/common/memory/byte_stream.h>

namespace kumo::pollux {
    StreamArena::StreamArena(memory::MemoryPool *pool) : pool_(pool) {
    }

    void StreamArena::newRange(
        int32_t bytes,
        ByteRange * /*lastRange*/,
        ByteRange *range) {
        POLLUX_CHECK_GT(bytes, 0, "StreamArena::newRange can't be zero length");
        const memory::MachinePageCount numPages =
                memory::AllocationTraits::numPages(bytes);
        // If new range 'bytes' is larger than the largest class page size, then we
        // allocate a large chunk of contiguous memory for this range.
        if (numPages > pool_->largestSizeClass()) {
            memory::ContiguousAllocation largeAllocation;
            pool_->allocateContiguous(numPages, largeAllocation);
            range->buffer = largeAllocation.data();
            range->size = largeAllocation.size();
            range->position = 0;
            size_ += range->size;
            largeAllocations_.push_back(std::move(largeAllocation));
            return;
        }

        const int32_t numRuns = allocation_.numRuns();
        if (currentRun_ >= numRuns) {
            if (numRuns > 0) {
                allocations_.push_back(
                    std::make_unique<memory::Allocation>(std::move(allocation_)));
            }
            pool_->allocateNonContiguous(
                std::max(allocationQuantum_, numPages), allocation_);
            currentRun_ = 0;
            currentOffset_ = 0;
            size_ += allocation_.byteSize();
        }
        auto run = allocation_.runAt(currentRun_);
        range->buffer = run.data() + currentOffset_;
        const int32_t availableBytes = run.numBytes() - currentOffset_;
        range->size = std::min(bytes, availableBytes);
        range->position = 0;
        currentOffset_ += range->size;
        POLLUX_DCHECK_LE(currentOffset_, run.numBytes());
        if (currentOffset_ == run.numBytes()) {
            ++currentRun_;
            ++currentOffset_ = 0;
        }
    }

    void StreamArena::clear() {
        allocations_.clear();
        pool_->freeNonContiguous(allocation_);
        currentRun_ = 0;
        currentOffset_ = 0;
        largeAllocations_.clear();
        size_ = 0;
    }
} // namespace kumo::pollux
